"""
Author: FM team
Date: 2023-4-28
"""

import os
from tqdm import tqdm

from wputils.utils.utils import generate_series_info_sorted_by_z_postion_asc
from wputils.utils.io import fpath, fdirs, wnii, wseq, rnii, gsiz, gmdf
from wputils.utils.resize import m3d
from wputils.utils.dicom2nii import get_series_files, SortSeriesFiles
from wputils.utils.utils import pe


def replace_blanks_with_underscore(text):
    """
    Replace all whitespace characters in a given string with underscores.

    Args:
        text (str): The input string to process.

    Returns:
        str: The modified string with all whitespace characters replaced by underscores.
    """
    return '_'.join(text.split())


def acquire_img(root):
    fs = [os.path.join(root, f) for f in os.listdir(root) if f.endswith('.nii.gz')]
    if len(fs) == 1:
        return fs[0]
    else:
        sizs = [gsiz(f) for f in fs]
        return fs[sizs.index(max(sizs))]


if __name__ == '__main__':

    # DCM_DIRS = '/share_data/mmd/dcm/'
    # DCM_FILES = fdirs(DCM_DIRS, '.dcm')
    #
    # failslist = []
    #
    # for dcm_file in tqdm(DCM_FILES):
    #
    #     try:
    #
    #         series_uid_to_files = get_series_files(dcm_file)
    #         # print(series_uid_to_files.keys())
    #         nifti_file = dcm_file.replace('/dcm/', '/nifti/')
    #
    #         for k, v in series_uid_to_files.items():
    #             if len(v) < 12:
    #                 continue
    #
    #             try:
    #                 ssf = SortSeriesFiles(v)
    #                 series_info, img_sitk = ssf.get_series_info_and_sitk()
    #                 series_description = series_info.get('series_description')
    #                 series_description = replace_blanks_with_underscore(
    #                     series_description) if series_description != '' else str(k)
    #                 nifti_path = os.path.join(nifti_file, series_description)
    #
    #                 wnii(img_sitk, file_path=os.path.join(nifti_path, 'img.nii.gz'))
    #                 wseq(series_info, os.path.join(nifti_path, 'info.json'))
    #
    #             except Exception as e:
    #                 print(f'get_series_info_and_sitk error: {dcm_file}')
    #                 print(pe(e))
    #                 failslist.append(dcm_file)
    #                 continue
    #
    #     except Exception as e:
    #         print(f'get_series_files error: {dcm_file}')
    #         print(pe(e))
    #         failslist.append(dcm_file)
    #         continue
    #
    # wseq(failslist, '/share_data/mmd/failslist.json')

    # ----------------------------------------

    NIFTI_ROOT = '/share_data/mmd/nifti'
    NIFTI_DIRS = fdirs(NIFTI_ROOT, '.nii.gz')
    SAVE_ROOT = '/share_data/mmd/nifti_resized'

    print(f'total: {len(NIFTI_DIRS)}, with example: {NIFTI_DIRS[0]}')

    for nifti in tqdm(NIFTI_DIRS):
        fimg = acquire_img(nifti)
        arr, stk = rnii(fimg)
        dims = arr.shape
        if 16 < dims[0] < 32:
            mdf = gmdf(fimg)
            parr = m3d(arr, (24, 512, 512))
            wnii(parr, stk, os.path.join(SAVE_ROOT, mdf + '.nii.gz'))
